import torch
from torch import nn
from src.model.feature_model.feature_block import FeatureBlock
from src.model.space_model.space_block import SpaceBlock
from src.model.time_model.time_block import TimeBlock
from src.model.dense import Dense
from src.model.time_embedding import Time_Embedding
from typing import Tuple


class Model(nn.Module):
    """模型
    """
    def __init__(self, 
                 target_turbine, 
                 target_feature, 
                 num_nodes, 
                 num_layer, 
                 seq, 
                 pred_len, 
                 dims: Tuple[int, int, int], 
                 static_graph=False, 
                 feature_block=True, 
                 only_time_block=False, 
                 single_point=True):
        super(Model, self).__init__()
        self.target_turbine = target_turbine
        self.target_feature = target_feature
        self.num_nodes = num_nodes
        self.num_layer = num_layer
        self.seq = seq
        self.pred_len = pred_len
        self.input_size, self.hidden_size, self.output_size = dims
        self.static_graph = static_graph
        self.use_feature_block = feature_block
        self.only_time_block = only_time_block
        self.single_point = single_point
        
        # 特征映射层
        dense_dims = (self.input_size, self.hidden_size, self.hidden_size // 2)
        self.dense = Dense(dense_dims)
        self.time_embedding = Time_Embedding(self.hidden_size // 2, self.seq)

        # 初始化特征模块
        feature_dims = (self.seq, 6, -1)
        self.featureblock = FeatureBlock(feature_dims)

        # 初始化时间模块
        time_dims = (self.hidden_size // 2, self.hidden_size, self.hidden_size)
        self.timeblock = TimeBlock(time_dims, self.seq, self.pred_len, self.num_nodes, self.num_layer, self.single_point)

        # 初始化空间模块
        space_dims = (self.hidden_size, self.hidden_size, self.hidden_size)
        self.spaceblock = SpaceBlock(space_dims, self.seq, self.num_nodes, self.num_layer, self.static_graph, self.single_point)

        # 初始化输出模块
        self.fc = nn.Linear(self.hidden_size, self.output_size)


    def forward(self, x, adj, is_training=True):
        x = self.dense(x)
        x = x.permute(0, 2, 3, 1)
        x = self.time_embedding(x)
        x = x.permute(0, 3, 1, 2).contiguous()
        
        if self.use_feature_block:
            x, a = self.featureblock(x, is_training)
        # y.shape = (B, P, N, H)，两个都是，如果是单点预测，y.shape = (B, N, H)
        if not self.only_time_block:
            yp = self.timeblock(x)
            #yp.shape = (B, N, H)，x.shape = (B, T, N, H)
            if is_training:
                y, adj = self.spaceblock(x, yp, adj, is_training)
            else:
                y, adj_g, adj_a, adj_d, adj_s = self.spaceblock(x, yp, adj, is_training)
            # 输出层将特征H压缩到输出的维度self.output_size
            if self.single_point:
                out = self.fc(y[:, self.target_turbine, :])
            else:
                out = self.fc(y[:, :, self.target_turbine, :])
        else:
            y = self.timeblock.gru_list[self.target_turbine](x[:, :, self.target_turbine, :])
            out = self.fc(y[:, -1, :])
        if is_training:
            return out.squeeze(-1)
        else:
            return out.squeeze(-1), a, adj_g, adj_a, adj_d, adj_s
